package sd

import (
	"fmt"
	"github.com/go-redis/redis/v8"
	"github.com/imroc/req/v3"
	"gorm.io/gorm"
	"myai/core/types"
	logger2 "myai/logger"
	"myai/service"
	"myai/service/oss"
	"myai/store"
	"myai/store/model"
	"myai/utils"
	"time"
)

var log = logger2.GetLogger()

type Service struct {
	httpClient  *req.Client
	taskQueue   *store.RedisQueue
	notifyQueue *store.RedisQueue
	db          *gorm.DB
	uploader    *oss.UploaderManager
	wsService   *service.WebSocketService
	userService *service.UserService
}

func NewService(db *gorm.DB, manager *oss.UploaderManager,
	redisCli *redis.Client, wsService *service.WebSocketService,
	userService *service.UserService) *Service {
	return &Service{
		httpClient:  req.C(),
		taskQueue:   store.NewRedisQueue("My_StableDiffusion_Task_Queue", redisCli),
		notifyQueue: store.NewRedisQueue("My_StableDiffusion_Notify_Queue", redisCli),
		db:          db,
		wsService:   wsService,
		uploader:    manager,
		userService: userService,
	}
}

// Txt2ImgReq Sd文生图请求实体
type Txt2ImgReq struct {
	Prompt            string  `json:"prompt"`                         // 主提示词
	NegativePrompt    string  `json:"negative_prompt"`                // 负面提示词
	Seed              int64   `json:"seed,omitempty"`                 // 随机种子，-1为随机
	Steps             int     `json:"steps"`                          // 采样步数
	CfgScale          float32 `json:"cfg_scale"`                      // 提示词引导强度
	Width             int     `json:"width"`                          // 图像宽度（像素）
	Height            int     `json:"height"`                         // 图像高度（像素）
	SamplerName       string  `json:"sampler_name"`                   // 采样器名称
	Scheduler         string  `json:"scheduler"`                      // 调度器类型
	EnableHr          bool    `json:"enable_hr,omitempty"`            // 是否启用高分辨率修复
	HrScale           int     `json:"hr_scale,omitempty"`             // 高分辨率缩放因子
	HrUpscaler        string  `json:"hr_upscaler,omitempty"`          // 高分辨率上采样器
	HrSecondPassSteps int     `json:"hr_second_pass_steps,omitempty"` // 高分辨率第二阶段步数
	DenoisingStrength float32 `json:"denoising_strength,omitempty"`   // 去噪强度（0-1）
	ForceTaskId       string  `json:"force_task_id,omitempty"`        // 强制任务Id
}

// Txt2ImgResp 文生图响应
type Txt2ImgResp struct {
	Images     []string `json:"images"`     // Base64 编码的生成图像
	Parameters struct{} `json:"parameters"` // 请求参数回显
	Info       string   `json:"info"`       // 包含模型版本、耗时等元数据
}

// TaskProgressResp 任务进度响应实体
type TaskProgressResp struct {
	Progress    float64 `json:"progress"`     // 当前进度百分比
	EtaRelative float64 `json:"eta_relative"` // 预估剩余时间（秒）
}

func (s *Service) Run() {
	// 把数据库中未提交的任务加载到队列中
	var jobs []model.SdJob
	s.db.Where("progress", 0).Find(&jobs)
	for _, job := range jobs {
		var task types.SdTask
		err := utils.JsonDecode(job.TaskInfo, &task)
		if err != nil {
			log.Errorf("解码SD任务信息出现错误：%v", err)
			continue
		}
		task.Id = int(job.Id)
		// 把任务发送到队列中
		s.PushTask(task)
	}
	log.Info("启动 Stable-Diffusion 任务消费者")
	go func() {
		for {
			var task types.SdTask
			err := s.taskQueue.LPop(&task)
			if err != nil {
				log.Errorf("SD 消费队列任务出现错误：%v", err)
				continue
			}
			// 翻译提示词
			if utils.HasChinese(task.Params.Prompt) {
				content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate,
					task.Params.Prompt), task.TranslateModelId)
				if err == nil {
					task.Params.Prompt = content
				} else {
					log.Warnf("SD 翻译提示词出现错误：%v", err)
				}
			}

			// 翻译反向提示词
			if task.Params.NegPrompt != "" && utils.HasChinese(task.Params.NegPrompt) {
				content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate,
					task.Params.NegPrompt), task.TranslateModelId)
				if err == nil {
					task.Params.NegPrompt = content
				} else {
					log.Warnf("SD 翻译反向提示词出现错误：%v", err)
				}
			}
			log.Infof("开始处理新的Stable-Diffusion任务：%+v", task)
			// 执行任务
			err = s.Txt2Img(task)
			if err != nil {
				log.Errorf("SD绘画任务执行失败：%v", err.Error())
				// 修改任务信息
				s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumns(map[string]interface{}{
					"progress": service.FailTaskProgress,
					"err_msg":  err.Error(),
				})
				// 通知前端，任务失败
				s.notifyQueue.RPush(service.NotifyMessage{
					UserId:   task.UserId,
					ClientId: task.ClientId,
					JobId:    task.Id,
					Message:  service.TaskStatusFailed,
				})
				continue
			}
		}
	}()
}

// Txt2Img 文生图API
func (s *Service) Txt2Img(task types.SdTask) error {
	body := Txt2ImgReq{
		Prompt:         task.Params.Prompt,
		NegativePrompt: task.Params.NegPrompt,
		Steps:          task.Params.Steps,
		CfgScale:       task.Params.CfgScale,
		Width:          task.Params.Width,
		Height:         task.Params.Height,
		SamplerName:    task.Params.Sampler,
		Scheduler:      task.Params.Scheduler,
		ForceTaskId:    task.Params.TaskId,
	}
	// 判断随机数种子是否大于0
	if task.Params.Seed > 0 {
		body.Seed = task.Params.Seed
	}
	// 启用高清修护
	if task.Params.HdFix {
		body.EnableHr = true
		body.HrScale = task.Params.HdScale
		body.HrUpscaler = task.Params.HdScaleAlg
		body.HrSecondPassSteps = task.Params.HdSteps
		body.DenoisingStrength = task.Params.HdRedrawRate
	}

	var res Txt2ImgResp
	var errChan = make(chan error)

	var apiKey model.ApiKey
	// 查询 SD 的 API 秘钥
	err := s.db.Where("type", "sd").Where("enabled", true).Order("last_used_at asc").First(&apiKey).Error
	if err != nil {
		return fmt.Errorf("没有可用的 Stable-Diffusion api 密钥：%v", err)
	}

	apiURL := fmt.Sprintf(service.SdV1Txt2imgPath, apiKey.ApiURL)
	log.Infof("向 SD 发送文生图请求：%s", apiURL)
	// 开始向 sd-api 端点发送请求
	go func() {
		// 发送请求
		response, err := s.httpClient.R().
			SetHeader("Authorization", apiKey.Value).
			SetBody(body).
			SetSuccessResult(&res).
			Post(apiURL)
		if err != nil {
			errChan <- err
			return
		}
		if response.IsErrorState() {
			errChan <- fmt.Errorf("错误的HTTP状态码：%v", response.Status)
			return
		}
		// 更新 API KEY 最后一次使用时间
		apiKey.LastUsedAt = time.Now().Unix()
		s.db.Updates(&apiKey)
		// 保存 Base64 图片
		imgURL, err := s.uploader.GetUploadHandler().PutBase64(res.Images[0])
		if err != nil {
			errChan <- fmt.Errorf("SD任务上传图像时出错：%v，图像Base64：%s", err, res.Images[0])
			return
		}
		// 获取绘画真实的 seed
		var info map[string]interface{}
		err = utils.JsonDecode(res.Info, &info)
		if err != nil {
			errChan <- fmt.Errorf("SD解码响应数据错误：%v", err)
			return
		}
		task.Params.Seed = int64(utils.IntValue(utils.InterfaceToString(info["seed"]), -1))
		s.db.Model(&model.SdJob{Id: uint(task.Id)}).
			UpdateColumns(model.SdJob{ImgURL: imgURL, Params: utils.JsonEncode(task.Params), Prompt: task.Params.Prompt})
		errChan <- nil
	}()
	// 等待任务完成
	for {
		select {
		case err := <-errChan:
			// err 不为空表示任务执行失败
			if err != nil {
				return err
			}
			// 任务完成，更新数据库任务进度
			s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", 100)
			// 发送通知
			s.notifyQueue.RPush(service.NotifyMessage{
				ClientId: task.ClientId,
				UserId:   task.UserId,
				JobId:    task.Id,
				Message:  service.TaskStatusFinished,
			})
			return nil
		default:
			// 获取任务进度
			err, resp := s.checkTaskProgress(apiKey)
			// 更新任务进度
			if err == nil && resp.Progress > 0 {
				s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", int(resp.Progress*100))
				// 发送更新状态信号，通知前端更新进度
				s.notifyQueue.RPush(service.NotifyMessage{
					ClientId: task.ClientId,
					UserId:   task.UserId,
					JobId:    task.Id,
					Message:  service.TaskStatusRunning,
				})
			}
			time.Sleep(time.Second)
		}
	}
	return nil
}

// checkTaskProgress 检查 Stable-Diffusion 任务进度
func (s *Service) checkTaskProgress(apiKey model.ApiKey) (error, *TaskProgressResp) {
	apiURL := fmt.Sprintf(service.SdV1ProgressPath, apiKey.ApiURL)
	var res TaskProgressResp
	response, err := s.httpClient.R().
		SetHeader("Authorization", apiKey.Value).
		SetSuccessResult(&res).
		Get(apiURL)
	if err != nil {
		return err, nil
	}
	if response.IsErrorState() {
		return fmt.Errorf("错误的HTTP状态码: %v", response.Status), nil
	}
	return nil, &res
}

// PushTask 添加新的Stable-Diffusion任务
func (s *Service) PushTask(task types.SdTask) {
	log.Debugf("将新的Stable-Diffusion任务添加到任务列表中：%+v", task)
	s.taskQueue.RPush(task)
}

// CheckTaskNotify 检查任务通知队列
func (s *Service) CheckTaskNotify() {
	go func() {
		log.Info("运行 Stable-Diffusion 任务通知检查 ...")
		for {
			var message service.NotifyMessage
			// 从队列中消费数据
			err := s.notifyQueue.LPop(&message)
			if err != nil {
				log.Errorf("SD 任务通知队列消费消息失败：%v", err.Error())
				continue
			}
			log.Debugf("SD任务通知消息：%+v", message)
			// 根据客户端Id获取一个客户端连接
			client := s.wsService.Clients.Get(message.ClientId)
			if client == nil {
				log.Warnf("客户端连接为空，客户端Id为：%s", message.ClientId)
				continue
			}
			// 向客户端用户发送消息
			utils.SendChannelMsg(client, types.ChSd, message.Message)
		}
	}()
}

// CheckTaskStatus 检查任务状态，自动删除过期或者失败的任务
func (s *Service) CheckTaskStatus() {
	go func() {
		log.Infof("运行 Stable-Diffusion 任务状态检查 ...")
		for {
			var jobs []model.SdJob
			res := s.db.Where("progress < ?", 100).Find(&jobs)
			if res.Error != nil {
				log.Errorf("查询正在执行的 SD JOB 任务出现错误：%v", res.Error.Error())
				time.Sleep(5 * time.Second)
				continue
			}
			for _, job := range jobs {
				// 五分钟还没完成标记任务失败
				if time.Now().Sub(job.CreatedAt) > time.Minute*5 {
					job.Progress = service.FailTaskProgress
					job.ErrMsg = "任务超时"
					s.db.Updates(&job)
				}
			}

			// 找出失败的任务，并恢复其扣减的算力
			res = s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
			if res.Error != nil {
				log.Errorf("查询失败的 SD JOB 任务出现错误：%v", res.Error.Error())
				continue
			}
			for _, job := range jobs {
				// 增加算力
				err := s.userService.IncreasePower(int(job.UserId), job.Power, model.PowerLog{
					Type:   types.PowerRefund,
					Model:  "stable-diffusion",
					Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d，Err：%s", job.Id, job.ErrMsg),
				})
				if err != nil {
					log.Errorf("在执行恢复扣减的算力时操作失败：%v", err.Error())
					continue
				}
				// 更新任务状态
				s.db.Model(&job).UpdateColumn("power", 0)
			}
			time.Sleep(time.Second * 5)
		}
	}()
}
